package com.konoha.controller;

import com.baidu.aip.ocr.AipOcr;
import com.konoha.pojo.RecognitionResult;
import com.konoha.pojo.Result;
import lombok.extern.slf4j.Slf4j;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.HashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@RestController
@Slf4j
@RequestMapping("/api")
public class BaiduOcrController {

    @Autowired
    private AipOcr aipOcr;

    @PostMapping(value = "/process", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public Result<RecognitionResult> processImage(@RequestParam("file") MultipartFile file) {
        try {
            byte[] imageBytes = file.getBytes();
            JSONObject result = aipOcr.basicGeneral(imageBytes, getOptions());

            // 解析百度OCR的识别结果
            RecognitionResult recognitionResult = parseOCRResult(result);
            return Result.success(recognitionResult);
        } catch (IOException e) {
            log.error("Error processing image: {}", e.getMessage());
            return Result.error("Error processing image");
        } catch (Exception e) {
            log.error("Error: {}", e.getMessage());
            return Result.error("Error");
        }
    }

    // 解析百度OCR的识别结果
    private RecognitionResult parseOCRResult(JSONObject result) throws JSONException {
        RecognitionResult recognitionResult = new RecognitionResult();

        // 检查"words_result"字段是否存在且是JSON数组
        if (result.has("words_result") && result.get("words_result") instanceof JSONArray) {
            JSONArray wordsResultArray = result.getJSONArray("words_result");

            // 定义标志位，表示是否找到了学号、姓名、成绩和作答用时
            boolean foundStudentId = false;
            boolean foundName = false;
            boolean foundScore = false;
            boolean foundTimeTaken = false;

            // 遍历识别结果
            for (int i = 0; i < wordsResultArray.length(); i++) {
                Object item = wordsResultArray.get(i);
                // 检查数组元素是否是JSONObject类型
                if (item instanceof JSONObject) {
                    JSONObject word = (JSONObject) item;
                    String text = word.getString("words");

                    // 判断是否包含学号、姓名、成绩和作答用时
                    if (text.contains("学号")) {
                        recognitionResult.setStudentId(removeColon(text.substring(text.indexOf("学号") + 2)));
                        foundStudentId = true;
                    } else if (text.contains("姓名")) {
                        recognitionResult.setName(removeColon(text.substring(text.indexOf("姓名") + 2)));
                        foundName = true;
                    } else if (text.contains("成绩")) {
                        // 从包含"成绩"关键字的行开始，向后搜索提取成绩
                        foundScore = true;
                        String score = extractScoreFromLine(wordsResultArray, i);
                        // 去除成绩中的作答用时部分
                        score = removeTimeTaken(score);
                        recognitionResult.setScore(score);
                    } else if (foundScore && text.matches(".*分钟.*")) {
                        // 如果已经找到成绩，并且当前文本包含"分钟"字样，则认为是作答用时
                        foundTimeTaken = true;
                        String timeTaken = extractTimeTaken(text);
                        recognitionResult.setTimeTaken(timeTaken);
                        break; // 找到作答用时后退出循环
                    }
                }
            }

            // 如果没有找到学号、姓名、成绩或作答用时，可以设置一个默认值或者给出提示
            if (!foundStudentId) {
                recognitionResult.setStudentId("未识别到学号");
            }
            if (!foundName) {
                recognitionResult.setName("未识别到姓名");
            }
            if (!foundScore) {
                recognitionResult.setScore("未识别到成绩");
            }
            if (!foundTimeTaken) {
                recognitionResult.setTimeTaken("未识别到作答用时");
            }
        }

        return recognitionResult;
    }

    // 从包含"成绩"关键字的行开始，向后搜索提取成绩
    private String extractScoreFromLine(JSONArray wordsResultArray, int startIndex) throws JSONException {
        StringBuilder scoreBuilder = new StringBuilder();
        // 从包含"成绩"关键字的行开始，向后搜索提取成绩
        for (int i = startIndex + 1; i < wordsResultArray.length(); i++) {
            Object item = wordsResultArray.get(i);
            if (item instanceof JSONObject) {
                JSONObject word = (JSONObject) item;
                String text = word.getString("words");
                // 判断当前行是否包含数字和"分"字样，如果是则将其添加到成绩字符串中
                if (text.matches(".*\\d+分.*")) {
                    scoreBuilder.append(text).append(" ");
                }
            }
        }
        return scoreBuilder.toString().trim();
    }

    // 提取作答用时（分钟）
    private String extractTimeTaken(String text) {
        // 使用正则表达式匹配分钟数，例如：94分钟
        Pattern pattern = Pattern.compile("\\b\\d+分钟\\b");
        Matcher matcher = pattern.matcher(text);
        if (matcher.find()) {
            return matcher.group();
        } else {
            return null;
        }
    }

    // 去除成绩中的作答用时部分
    private String removeTimeTaken(String score) {
        // 使用正则表达式去除成绩中的作答用时部分
        return score.replaceAll("\\s*\\d+分钟\\s*", "");
    }

    // 去除冒号
    private String removeColon(String text) {
        return text.replaceAll("：", "");
    }

    // 获取OCR识别选项
    private HashMap<String, String> getOptions() {
        HashMap<String, String> options = new HashMap<>();
        options.put("detect_direction", "true");
        options.put("probability", "true");
        return options;
    }
}
